import asyncio
import logging
from typing import List, Optional, Set
from dataclasses import dataclass
import time
from src.config.settings import K_BUCKET_SIZE, NODE_ID_LENGTH, NODE_TIMEOUT,NODE_PING_COUNT,SAVE_DB
from src.log.logger import setup_logger

@dataclass
class Node:
    """DHT节点信息"""
    node_id: bytes
    ip: str
    port: int
    last_seen: float = 0.0
    ping_count: int = 0
    execute_find_node: bool = False
    
    _lock_ping_count: asyncio.Lock = asyncio.Lock()
    _lock_last_seen: asyncio.Lock = asyncio.Lock()
    _lock_execute_find_node: asyncio.Lock = asyncio.Lock()

    def __hash__(self):
        return hash((self.node_id, self.ip, self.port))

    def __eq__(self, other):
        if not isinstance(other, Node):
            return False
        return (self.node_id == other.node_id and 
                self.ip == other.ip and 
                self.port == other.port)
    
    async def set_execute_find_node(self,value:bool):
        async with self._lock_execute_find_node:
            self.execute_find_node = value

    async def update_last_seen(self):
        """更新节点最后可见时间"""
        async with self._lock_last_seen:
            self.last_seen = time.time()

    @property
    def is_active(self) -> bool:
        """检查节点是否活跃"""
        return time.time() - self.last_seen < NODE_TIMEOUT and self.ping_count > NODE_PING_COUNT
    
    async def ping_count_add(self):
        async with self._lock_ping_count:
            self.ping_count += 1
    
    async def change_execute_find_node(self,value:bool):
        async with self._lock_execute_find_node:
            self.execute_find_node = value

class KBucket:
    """K桶实现"""
    def __init__(self, min_id: int, max_id: int):
        self.min_id = min_id  # 桶的最小ID范围
        self.max_id = max_id  # 桶的最大ID范围
        self.nodes: List[Node] = []  # 桶中的节点列表
        self.last_updated = time.time()

    def add_node(self, node: Node) -> bool:
        """添加节点到K桶
        
        Args:
            node: 要添加的节点
            
        Returns:
            bool: 是否添加成功
        """
        if len(self.nodes) < K_BUCKET_SIZE:
            if node not in self.nodes:
                self.nodes.append(node)
                self.last_updated = time.time()
                return True
        return False

    def remove_node(self, node: Node) -> None:
        """从K桶中移除节点"""
        if node in self.nodes:
            self.nodes.remove(node)
            self.last_updated = time.time()

    def get_nodes(self) -> List[Node]:
        """获取桶中的所有活跃节点"""
        return [node for node in self.nodes if node.is_active]

class RoutingTable:
    """DHT路由表"""
    def __init__(self, node_id: bytes):
        self.node_id = node_id
        self.buckets: List[KBucket] = []
        self.logger: logging.Logger = setup_logger("routing_table")
        self.save_db:bool = SAVE_DB
        self._initialize_buckets()

    def _initialize_buckets(self):
        """初始化K桶"""
        for i in range(NODE_ID_LENGTH * 8):  # 对每个比特位创建一个K桶
            min_id = 2 ** i
            max_id = 2 ** (i + 1) - 1
            self.buckets.append(KBucket(min_id, max_id))

    def _get_bucket_index(self, node_id: bytes) -> int:
        """计算节点应该在哪个K桶中
        
        Args:
            node_id: 节点ID
            
        Returns:
            int: K桶索引
        """
        distance = int.from_bytes(node_id, 'big') ^ int.from_bytes(self.node_id, 'big')
        if distance == 0:
            return 0
        return (distance - 1).bit_length() - 1

    def add_node(self, node: Node) -> bool:
        """添加节点到路由表
        
        Args:
            node: 要添加的节点
            
        Returns:
            bool: 是否添加成功
        """
        if node.node_id == self.node_id:
            return False
        
        bucket_index = self._get_bucket_index(node.node_id)
        asyncio.create_task(node.update_last_seen())
        return  self.buckets[bucket_index].add_node(node)

    def remove_node(self, node: Node) -> None:
        """从路由表中移除节点"""
        bucket_index = self._get_bucket_index(node.node_id)
        self.buckets[bucket_index].remove_node(node)

    def get_closest_nodes(self, target_id: bytes, count: int = K_BUCKET_SIZE) -> List[Node]:
        """获取距离目标节点最近的K个节点
        
        Args:
            target_id: 目标节点ID
            count: 需要返回的节点数量
            
        Returns:
            List[Node]: 最近的节点列表
        """
        nodes: Set[Node] = set()
        target_bucket = self._get_bucket_index(target_id)
        
        # 从目标桶开始，向两边扩展搜索
        left = right = target_bucket
        while (left >= 0 or right < len(self.buckets)) and len(nodes) < count:
            if left >= 0:
                nodes.update(self.buckets[left].get_nodes())
                left -= 1
            if right < len(self.buckets):
                nodes.update(self.buckets[right].get_nodes())
                right += 1
                
        # 按照到目标节点的距离排序
        target_int = int.from_bytes(target_id, 'big')
        sorted_nodes = sorted(
            nodes,
            key=lambda x: int.from_bytes(x.node_id, 'big') ^ target_int
        )
        
        return sorted_nodes[:count]

    def get_all_nodes(self) -> List[Node]:
        """获取路由表中的所有活跃节点"""
        nodes = []
        for bucket in self.buckets:
            nodes.extend(bucket.get_nodes())
        return nodes 

    def has_node(self, node_id: bytes = None, ip: str = None, port: int = None) -> bool:
        """检查节点是否存在于路由表中
        
        可以通过节点ID或者IP+端口的组合来查找
        
        Args:
            node_id: 节点ID
            ip: 节点IP地址
            port: 节点端口
            
        Returns:
            bool: 节点是否存在
        """
        if not any([node_id, ip, port]):
            return False
            
        for bucket in self.buckets:
            for node in bucket.nodes:
                # 如果指定了节点ID，优先用节点ID匹配
                if node_id and node.node_id == node_id:
                    return True
                # 如果指定了IP和端口，用IP+端口匹配
                elif ip and port and node.ip == ip and node.port == port:
                    return True
        return False

    def find_node(self, node_id: bytes = None, ip: str = None, port: int = None) -> Optional[Node]:
        """查找指定节点
        
        可以通过节点ID或者IP+端口的组合来查找
        
        Args:
            node_id: 节点ID
            ip: 节点IP地址
            port: 节点端口
            
        Returns:
            Optional[Node]: 找到的节点，如果不存在则返回None
        """
        if not any([node_id, ip, port]):
            return None
            
        for bucket in self.buckets:
            for node in bucket.nodes:
                # 如果指定了节点ID，优先用节点ID匹配
                if node_id and node.node_id == node_id:
                    return node
                # 如果指定了IP和端口，用IP+端口匹配
                elif ip and port and node.ip == ip and node.port == port:
                    return node
        return None

    def get_node_bucket(self, node_id: bytes) -> Optional[KBucket]:
        """获取指定节点所在的K桶
        
        Args:
            node_id: 节点ID
            
        Returns:
            Optional[KBucket]: 节点所在的K桶，如果不存在则返回None
        """
        bucket_index = self._get_bucket_index(node_id)
        return self.buckets[bucket_index] 

    async def update_node(self, node: Node) -> bool:
        """更新节点信息，如果节点不存在则添加新节点
        
        Args:
            node: 要更新的节点
            
        Returns:
            bool: 更新/添加是否成功
        """
        # 先查找现有节点
        existing_node = self.find_node(node_id=node.node_id)
        if existing_node:
            # 更新现有节点信息
            existing_node.ip = node.ip
            existing_node.port = node.port
            await existing_node.update_last_seen()
            await existing_node.ping_count_add()  
            return True
        else:
            # 添加新节点
            return self.add_node(node)